package DailyQuestion;

import java.util.Arrays;
import java.util.HashSet;
import java.util.Set;
import java.util.Stack;

public class 子序列最大优雅度_2813 {
    public long findMaximumElegance(int[][] items, int k) {
        Arrays.sort(items, (a, b) -> b[0] - a[0]); // 从大到小排序
        long ans = 0, totalProfit = 0;
        Set<Integer> isVis = new HashSet<>(); // 检测category是否遇到过
        Stack<Integer> duplicate = new Stack<>();

        for (int i = 0; i < items.length; i++) {
            int profit = items[i][0];
            int category = items[i][1];
            if (i < k) {
                totalProfit += profit;
                if (!isVis.add(category)) { // 有重复类别, 利润入栈
                    duplicate.push(profit);
                }
            } else if (!duplicate.isEmpty() && isVis.add(category)) { // 没遇到过,并且栈里有重复的类别
                totalProfit += profit - duplicate.pop();
            } // else: 比前面利润小，并且类别还重复（-1+1），不考虑
            ans = Math.max(ans, totalProfit + (long) isVis.size() * isVis.size());
        }
        return ans;
    }
}
